// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <functional>
#include <tuple>
#include <vector>

#include <tbir/pipeline/attributes.h>
#include <tbir/pipeline/node.h>

namespace tbir::runtime {

    class TXSession;

    struct NativeOpKernel {
        virtual const char *ClassName() const = 0;

        virtual void Initialize(const String &resource_path, Attributes attrs) = 0;

        virtual RTValue Process(PyArgs input_args) const = 0;
    };

    struct NativeOp {
    public:
        NativeOp() = default;

        virtual ~NativeOp() = default;

        const std::string &InstanceName() const {
            return instance_name_;
        }

        const std::string &OpName() const {
            return instance_name_;
        }

        Symbol ComposeSymbol(const std::vector<const Symbol *> &args,
                             const ska::flat_hash_map<std::string, const Symbol *> &kwargs) const;

        Symbol ComposeSymbol(const std::vector<const Symbol *> &args) const {
            ska::flat_hash_map<std::string, const Symbol *> kwargs;
            return ComposeSymbol(args, kwargs);
        }

    protected:
        virtual NodePtr CreateNode() const;

    protected:
        TXSession *session_{nullptr};
        std::string op_name_;
        std::string instance_name_;
        Attributes generic_attrs_;
    };

}  // namespace tbir::runtime
